#pragma once

#include "AVLT.h"
#include <Windows.h>


struct TreeItemInfo
{
	std::string info{" "}; // 节点文本信息
	bool bIsDirty = true; // 节点颜色，true -> 黑色，false -> 红色
};

template<typename T>
struct RBNode: public NodeBase<T, RBNode<T>>
{
	RBNode():NodeBase<T, RBNode<T>>(), bIsDirty(false){}
	RBNode(const T& inData):NodeBase<T, RBNode<T>>(inData), bIsDirty(false){}
	
	// 获取祖父节点
	RBNode<T>* GrandParent() const
	{
		if(this->parent && this->parent->parent)
			return this->parent->parent;
		else
			return nullptr;
	}

	// 获取舅舅节点
	RBNode<T>* Uncle() const
	{
		// 如果祖父节点存在
		RBNode<T>* grandNode = GrandParent();
		if(grandNode)
		{
			// 如果父节点是祖父节点的左子节点
			// 舅舅节点就是祖父节点的右子节点。反之，是左子节点。
			if(this->parent == grandNode->left)
				return grandNode->right;
			else
				return grandNode->left;
		}
		else
			return nullptr;
	}

	// 获取兄弟节点
	RBNode<T>* Sibling() const
	{
		// 如果当前节点是父节点的左子节点，兄弟节点就是父节点的右子节点
		// 反之，兄弟节点就是父节点的左子节点
		if(this == this->parent->left)
			return this->parent->right;
		else
			return this->parent->left;
	}
	
	bool bIsDirty; // true 为黑，false 为红
};

// 红黑树
template<typename T, typename NodeType=RBNode<T>>
class RBT: public AVLT<T, NodeType>
{
public:
	using Node = NodeType;

	// 打印图形
	void ShowTreeGraph();
	
	// 插入
	void Insert(const T& element);

	// 删除
	void Delete(const T& element);
private:
	// 按情形插入
	void Insert_Case(Node* inNode);

	// 递归插入
	void InsertRecursively(Node*& parent, const T& element);

	// 左旋
	void LeftRotate(Node* inNode);

	// 右旋
	void RightRotate(Node* inNode);

	// 设置控制台颜色
	void SetConsoleColor(int inColor);

	/**
	 * 填充二叉树节点绘制缓冲区
	 * @param childRoot 
	 * @param buffer 
	 * @param row 缓冲区行
	 * @param col 缓冲区列
	 * @param hier 数的总层数
	 */
	void FillDrawBuffer(Node* childRoot, std::vector<std::vector<TreeItemInfo>>& buffer,
						int row, int col, int hier) const;

	// 删除元素
	void Delete_Child(Node* inNode, const T& val);
	// 删除子节点
	void Delete_One_Child(Node* inNode);
	// 分情况平衡红黑树
	void Delete_Case(Node* node);
};

template <typename T, typename NodeType>
void RBT<T, NodeType>::ShowTreeGraph()
{
	if(!this->root)
		return;

	// 矩阵行数 == 树层级数
	int matrixHeight = this->Height() + 1;
	// 矩阵列数 == (2 * 树最深层的最大节点数) - 1
	int matrixWidth = std::pow(2, matrixHeight) - 1;
	std::vector<std::vector<TreeItemInfo>> treeMatrix(matrixHeight, std::vector<TreeItemInfo>(matrixWidth));

	FillDrawBuffer(this->root, treeMatrix, 0, (matrixWidth - 1)/2, matrixHeight);

	for(const auto& row : treeMatrix)
	{
		for(const auto& item : row)
		{
			if(item.bIsDirty)
				SetConsoleColor(FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_BLUE);
			else
				SetConsoleColor(FOREGROUND_RED);
			
			std::cout << item.info ;
		}
		std::cout << std::endl;
	}
	SetConsoleColor(FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_BLUE);
	
	std::cout << "\n" << std::endl;
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::Insert(const T& element)
{
	if(this->root == nullptr)
	{
		// 如果根为空
		Node* temp = new Node(element);
		this->root = temp;
		Insert_Case(temp);
	}
	else
	{
		// 如果根不为空
		InsertRecursively(this->root, element);
	}
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::InsertRecursively(Node*& parent, const T& element)
{
	if(element < parent->data)
	{
		// 沿左子树移动
		if(parent->left == nullptr)
		{
			Node* temp = new Node(element);
			temp->parent = parent;
			parent->left = temp;
			Insert_Case(temp);
		}
		else
			InsertRecursively(parent->left, element);
	}
	else if(element > parent->data)
	{
		// 沿右子树移动
		if(parent->right == nullptr)
		{
			Node* temp = new Node(element);
			temp->parent = parent;
			parent->right = temp;
			Insert_Case(temp);
		}
		else
			InsertRecursively(parent->right, element);
	}
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::LeftRotate(Node* inNode)
{
	Node* grandParent = inNode->GrandParent();
	Node* parent = inNode->parent;
	Node* childL = inNode->left;

	parent->right = childL;
	if(childL)
		childL->parent = parent;

	inNode->left = parent;
	parent->parent = inNode;

	// 如果原父节点是根节点，将根指针指向新的根节点（左旋过后的根节点是原来的新节点）
	if(this->root == parent)
		this->root = inNode;

	// 新节点的父节点更新为原来的祖父节点
	inNode->parent = grandParent;

	// 如果祖父节点存在
	if(grandParent)
	{
		// 如果原父节点是祖父节点的左子节点，将祖父节点的左子节点更新为新节点
		// 否则更新右子节点
		if(parent == grandParent->left)
			grandParent->left = inNode;
		else
			grandParent->right = inNode;
	}
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::RightRotate(Node* inNode)
{
	Node* grandParent = inNode->GrandParent();
	Node* parent = inNode->parent;
	Node* childR = inNode->right;

	parent->left = childR;
	if(childR)
		childR->parent = parent;

	inNode->right = parent;
	parent->parent = inNode;

	// 如果原父节点是根节点，将根指针指向新的根节点（右旋过后的根节点是原来的新节点）
	if(this->root == parent)
		this->root = inNode;

	// 新节点的父节点更新为原来的祖父节点
	inNode->parent = grandParent;

	// 如果祖父节点存在
	if(grandParent)
	{
		// 如果原父节点是祖父节点的左子节点，将祖父节点的左子节点更新为新节点
		// 否则更新右子节点
		if(parent == grandParent->left)
			grandParent->left = inNode;
		else
			grandParent->right = inNode;
	}
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::SetConsoleColor(int inColor)
{
	HANDLE handle = GetStdHandle(STD_OUTPUT_HANDLE);
	SetConsoleTextAttribute(handle, inColor);
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::FillDrawBuffer(Node* childRoot, std::vector<std::vector<TreeItemInfo>>& buffer, int row, int col,
	int hier) const
{
	if(!childRoot)
		return;

	// 填充元素
	TreeItemInfo itemInfo;
	itemInfo.info = std::to_string(childRoot->data);
	itemInfo.bIsDirty = childRoot->bIsDirty;
	buffer[row][col] = itemInfo;

	// 计算下一层元素列位置偏移
	// hier - row == 当前层数
	// 当前层数 - 1 == 下一层数
	// 层数 - 1 = 当前层的实际树高
	// hier - row - 2 == 下一层的实际树高
	// gap == 下一层子树的最深层的最大节点数
	int gap = std::pow(2, hier - row - 2);

	if(childRoot->left)
		FillDrawBuffer(childRoot->left, buffer, row + 1, col - gap, hier);
	else if( row + 1 < hier)
		buffer[row+1][col - gap].info = ".";
	
	if(childRoot->right)
		FillDrawBuffer(childRoot->right, buffer, row + 1, col + gap, hier);
	else if( row + 1 < hier)
		buffer[row+1][col + gap].info = ".";
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::Insert_Case(Node* inNode)
{
	if(inNode->parent == nullptr)
	{
		// 如果父节点为空
		
		/*
		 * 情形1 -> 父节点为空，说明该节点是根节点，设置为黑色
		 */
		
		this->root->bIsDirty = true;
	}
	else
	{
		// 如果父节点不为空

		Node* parent = inNode->parent;
		Node* uncle = inNode->Uncle();
		Node* grandParent = inNode->GrandParent();

		auto InsertCase_5 = [this](Node* inNode)
		{
			Node* parent = inNode->parent;
			Node* grandParent = inNode->GrandParent();

			// 将父节点染为黑色
			parent->bIsDirty = true;
			// 将祖父节点染为红色
			grandParent->bIsDirty = false;

			if((inNode == parent->left) && (parent == grandParent->left))
			{
				/*
				 *	如果新节点是父节点的左子节点 并且 父节点是祖父节点的左子节点
				 *	对父节点右旋
				 */
				RightRotate(parent);
			}
			else if((inNode == parent->right) && (parent == grandParent->right))
			{
				/*
				 *	如果新节点是父节点的右子节点 并且 父节点是祖父节点的右子节点
				 *	对父节点左旋
				 */
				LeftRotate(parent);
			}
		};
		
		if(inNode->parent->bIsDirty)
		{
			/*
			 *	情形2 -> 父节点是黑色
			 *	此时，树仍旧有效。所以什么也不做。
			 */
			
			return;
		}
		else if((uncle != nullptr) && !uncle->bIsDirty)
		{
			/*
			 *	情形3 -> 父节点是红色，叔父节点存在，并且叔父节点也是红色 （如果叔父节点是黑色，就违反了性质5）
			 *	将父节点和叔父节点都设置为黑色，再将祖父节点设置为红色。（用来保持性质5）
			 *	但是，祖父节点的父节点有可能也是红色，这样就违反了性质4.
			 *	所以再递归祖父节点.
			 */
			
			inNode->parent->bIsDirty = true;
			uncle->bIsDirty = true;
			grandParent->bIsDirty = false;
			Insert_Case(grandParent);
		}
		else if((inNode == parent->right) && (parent == grandParent->left))
		{
			/*
			 *	情形4-1 -> 父节点是红色，叔父节点不存在
			 *	新节点是父节点的右子节点，并且 父节点是祖父节点的左子节点
			 */

			// 左旋
			LeftRotate(inNode);
			// 更新 inNode
			inNode = inNode->left;

			InsertCase_5(inNode);
		}
		else if((inNode == parent->left) && (parent == grandParent->right))
		{
			/*
			 *	情形4-2 -> 父节点是红色，叔父节点不存在
			 *	新节点是父节点的左子节点，并且 父节点是祖父节点的右子节点
			 */

			// 右旋
			RightRotate(inNode);
			// 更新 inNode
			inNode = inNode->right;

			InsertCase_5(inNode);
		}
		else
		{
			/*
			 *	情形5 -> 父节点是红色，叔父节点不存在
			 *	新节点是父节点的左子节点，父节点是祖父节点的左子节点 或
			 *	新节点是父节点的右子节点，父节点是祖父节点的右子节点
			 */
			InsertCase_5(inNode);
		}

	}
	
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::Delete(const T& element)
{
	if(this->root == nullptr)
		return;

	/*
	 *	红黑树节点的删除情形
	 *	
	 *	如果被删除节点有两个子节点，那么可以将其转换为删除拥有一个或没有子节点的节点的问题。
	 *	例如：如果被删除节点有两个子节点，和二叉树的删除类似。寻找右子树中的值最小的节点，之后替换当前被删除节点的值，然后去删除这个最小右子树节点。
	 *		而最小右子树节点最多只有一个右子节点。如果被删除节点只有左子树，就删除这个节点本身。
	 */
	Delete_Child(this->root, element);
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::Delete_Child(Node* inNode, const T& val)
{
	if(val < inNode->data)
	{
		// 删除元素值小于当前节点的元素值，递归左子树
		Delete_Child(inNode->left, val);
	}
	else if(val > inNode->data)
	{
		// 删除元素值大于当前节点的元素值，递归右子树
		Delete_Child(inNode->right, val);
	}
	else
	{
		/*
		 *	删除元素值等于当前节点的元素值
		 *	将删除任何可能拥有两个子节点的问题，转换为删除拥有一个或没有子节点的节点问题
		 */
		
		if(inNode->right == nullptr)
		{
			// 如果右子树不存在
			Delete_One_Child(inNode);
		}
		else
		{
			// 如果右子树存在
			// 找到右子树中值最小的节点
			// 右子树中的最小值节点，要么只有一个右子节点，要么子节点为空
			Node* minimum = this->FindMin(inNode->right);
			// 替换当前节点的值
			inNode->data = minimum->data;
			Delete_One_Child(minimum);
		}
	}
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::Delete_One_Child(Node* inNode)
{
	/*
	 *	这一部分处理的节点都只有一个或没有子节点
	 */
	
	// 如果当前节点的左子节点为空，获取右子节点，反之，获取左子节点
	Node* child = inNode->left == nullptr ? inNode->right : inNode->left;
	Node* parent = inNode->parent;
	
	/*
	 *	检查被删除节点是否为根节点
	 */
	if(parent == nullptr)
	{
		delete inNode;

		// 如果子节点存在
		if(child)
		{
			// 子节点将成为新的根节点
			child->parent = nullptr;
			this->root = child;
			// 将节点染为黑色
			this->root->bIsDirty = true;
		}
		else
			this->root = nullptr;

		return;
	}
	
	// 预平衡红黑树
	// 如果当前节点是黑色
	if(inNode->bIsDirty)
	{
		// 如果子节点存在，且为红色，将其染黑
		if(child && !child->bIsDirty)
			child->bIsDirty = true;
		// 若子节点不存在
		else
			Delete_Case(inNode);
	}
	
	// 用被删除节点的子节点替换该节点在树中的位置。
	if(inNode == parent->left)
	{
		// 如果当前节点是父节点的左子节点
		// 将父节点的左子节点替换为 child
		parent->left = child;
	}
	else
	{
		// 如果当前节点是父节点的右子节点
		// 将父节点的右子节点替换为 child
		parent->right = child;
	}
	// 如果子节点存在，更新子节点指向的父节点
	if(child)
		child->parent = parent;
	
	delete inNode;
}

template <typename T, typename NodeType>
void RBT<T, NodeType>::Delete_Case(Node* node)
{
	/*
	 *	这一部分处理的节点是黑色且没有子节点。
	 */
	// 父节点
	Node* parent = node->parent;
	// 兄弟节点
	Node* sibling = node->Sibling();

	// 交换两节点颜色
	auto SwapColor = [](Node* a, Node* b)
	{
		bool temp = b->bIsDirty;
		b->bIsDirty = a->bIsDirty;
		a->bIsDirty = temp;
	};

	// 如果兄弟节点是黑色
	if(sibling->bIsDirty)
	{
		/*
		 *	Case1 -> 如果兄弟节点是黑色，并且兄弟节点有一个与其方向一致的红色子节点（另一个方向不一致的节点可有可无）。
		 */
		if((sibling == sibling->parent->left) && sibling->left && !sibling->left->bIsDirty)
		{
			// 如果兄弟节点是其父节点的左子节点，兄弟节点的子节点是兄弟节点的左子节点，并且该子节点是红色

			// 交换父节点与兄弟节点的颜色
			SwapColor(sibling, sibling->parent);
			// 将兄弟节点的子节点染黑
			sibling->left->bIsDirty = true;
			// 以兄弟节点为中心，右旋
			RightRotate(sibling);
		}
		else if((sibling == sibling->parent->right) && sibling->right && !sibling->right->bIsDirty)
		{
			// 如果兄弟节点是其父节点的右子节点，兄弟节点的子节点是兄弟节点的右子节点，并且该子节点是红色

			// 交换父节点与兄弟节点的颜色
			SwapColor(sibling, sibling->parent);
			// 将兄弟节点的子节点染黑
			sibling->right->bIsDirty = true;
			// 以兄弟节点为中心，左旋
			LeftRotate(sibling);
		}
		/*
		 *	Case2 -> 如果兄弟节点是黑色，并且兄弟节点有一个方向与其不一致的红色子节点，另一个方向一致的子节点可有可无。
		 */
		else if((sibling == sibling->parent->left) && sibling->right && !sibling->right->bIsDirty)
		{
			// 如果兄弟节点是其父节点的左子节点，兄弟节点有一个右子节点，并且该子节点是红色

			// 交换子节点与兄弟节点的颜色
			SwapColor(sibling->right, sibling);
			// 以子节点为中心，左旋
			LeftRotate(sibling->right);
			// 递归调用，在 Case1 处理
			Delete_Case(node);
		}
		else if((sibling == sibling->parent->right) && sibling->left && !sibling->left->bIsDirty)
		{
			// 如果兄弟节点是其父节点的右子节点，兄弟节点有一个左子节点，并且该子节点是红色

			// 交换子节点与兄弟节点的颜色
			SwapColor(sibling->left, sibling);
			// 以子节点为中心，右旋
			RightRotate(sibling->left);
			// 递归调用，在 Case1 处理
			Delete_Case(node);
		}
		/*
		 *	Case3 -> 如果兄弟节点是黑色，兄弟节点没有子节点，父节点是红色或者黑色
		 */
		else if(!sibling->left && !sibling->right)
		{
			if(!parent->bIsDirty)
			{
				// 如果父节点是红色

				// 交换父节点和兄弟节点颜色
				SwapColor(parent, sibling);
			}
			else
			{
				// 如果父节点是黑色

				// 将兄弟节点染红，此时达到局部平衡
				sibling->bIsDirty = false;
				Delete_Case(parent);
			}
		}
	}
	else
	{
		/*
		 *	Case4 -> 如果兄弟节点是红色，此时兄弟节点必然有两个黑色子节点，且父节点是黑色，否则违反性质4、5.
		 */
		if(sibling == parent->left)
		{
			// 如果兄弟节点是父节点的左子节点

			// 交换父节点和兄弟节点的颜色
			SwapColor(sibling, parent);
			// 以兄弟节点为中心，右旋
			RightRotate(sibling);
		}
		else
		{
			// 如果兄弟节点是父节点的右子节点

			// 交换父节点和兄弟节点的颜色
			SwapColor(sibling, parent);
			// 以兄弟节点为中心，左旋
			LeftRotate(sibling);
		}
		// 递归处理
		Delete_Case(node);
	}
		
}
